#include "copycat.h"
static CUProgress progress;


//////////////////////////////////////////////////////////////////////////////
bool CCopycatTrain::Train( char *copyopt_str, char *parse_fn, char *giza_fn, char *align_fn, char *out_fn )
{
//	copyopt.depth=2;

	if( LoadCorpus( parse_fn, giza_fn, align_fn ) == false )	return false;


	if( out_files.Prepare( out_fn ) == false ) 	return false;
	if( ExtractSrcSubtree() == false )	return false;

	CUString optstr = copyopt_str;
	optstr.MakeToUpper();

#ifdef DEBUG
	fprintf( stdout, "Copyopt: %s\n", optstr.GetStr() ); fflush( stdout );
#endif

	if( optstr.Count( "R" ) == 1 )	{
		copyopt.enable_reorder = true;
		//	if( ExtractRuleCross() == false )	return false;
		if( ExtractRuleRF() == false )	return false;
	}
	if( optstr.Count( "I" ) == 1 )	{
		copyopt.enable_insert = true;
		if( ExtractInsRule() == false )	return false;
	}
	if( optstr.Count( "D" ) == 1 )	{
		copyopt.enable_delete = true;
	}

	return true;
}


//////////////////////////////////////////////////////////////////////////////
bool CCopycatTrain::ExtractSrcSubtree()
{
#ifdef DEBUG
	fprintf( stdout, "ExtractSrcSubtree start \n" ); fflush( stdout );
#endif

	stree.Clear();

	for( int i=0; i<num_sent; i++ )	{
		fprintf( stderr, "\r[Extract Src Subtree]   %d / %d sentences processed (# Subtree: %d)", i, num_sent, stree.GetSize() );

		CSent &s = sent[i];


		// load parse tree
		CTree t ;
		t.Init();
		if( t.LoadTree( s.parse ) == false )	{
			continue;
			//return false;
		}
		if( SetNodeSpan( s, t ) == false )	{
			fprintf( stderr, "set node span error %d\n", i );
			return false;
		}

		// for each node
		for( int j=0; j<t.num_node; j++ )	{
			CNode &n = t.node[j];

			if( MCH == "2" )	{ if( n.num_child != 2 )	continue; }
			else if( MCH == "23" )	 { if( n.num_child == 2  || n.num_child ==3 )	{ } else	{ continue; } }
			else { if( n.num_child <= 1 ) continue;  }


			if( IsAllChildrenAligned( t, n ) == false )	continue;

			// extract src patterns
			int num_cross = GetNumCross( t, n, s );
			CUString rule_str = GetSubtreeStr( t, n );
			stree.Add( rule_str, num_cross );
		}

	}
	fprintf( stderr, "\r[Extract Src Subtree]   %d / %d sentences processed (# Subtree: %d)\n",
					num_sent, num_sent, stree.GetSize() );

	if( stree.SaveToFile( out_files.file_src_info.GetFP() ) == false )	{
		fprintf( stderr, "[Fail to Save]\n" );
		return false;	
	}

#ifdef DEBUG
	fprintf( stdout, "ExtractSrcSubtree finish\n" ); fflush( stdout );
#endif
	return true;	
}


//////////////////////////////////////////////////////////////////////////////
void CCopycatTrain::RecoverReorder( CNode &r_node, CNode &org_node )
{
	for( int i=0; i<r_node.num_child; i++ )	{
		r_node.child[i] = org_node.child[i];
	}

}


//////////////////////////////////////////////////////////////////////////////
bool CCopycatTrain::ExtractRuleCross()
{
	for( int i=0; i<num_sent; i++ )	{
		fprintf( stderr, "\r[Extract Rule: CR   ]   %d / %d sentences processed (# Rules: %d)", i, num_sent, rule_cross.GetNumRules() );

		CSent &s = sent[i];

		// load tree
		CTree t, t2;
		t.Init(); t2.Init();
		if( t.LoadTree( s.parse ) == false )	{
			continue;
		//	return false;
		}
		if( t2.LoadTree( s.parse ) == false )	{
			continue;
		//	return false;
		}
		SetNodeSpan( s, t );
		SetNodeSpan( s, t2 );

		// store src statics
		for( int j=0; j<t.num_node; j++ )	{
			CNode &n = t.node[j];
			CNode &n2 = t2.node[j];

			//if( n.num_child != 2 )	continue;
			if( n.num_child == 2 )	{
//				if( IsAllChildrenAligned( t, n ) == false )	continue;

				CRuleCross new_rule;
				new_rule.before = GetSubtreeStr( t, n );
				new_rule.before_cross = GetNumCross( t2, n2,s );
				new_rule.freq = 1;

				// reorder
				int tmp = n2.child[0];
				n2.child[0] = n2.child[1];
				n2.child[1] = tmp;

				// insert
				new_rule.after = "A1 A0";
				new_rule.after_cross = GetNumCross( t2, n2, s );
				rule_cross.Add( new_rule );
				RecoverReorder( n2, n );
			}
			else if( n.num_child == 3 )	{
				continue;
//				if( IsAllChildrenAligned( t, n ) == false )	continue;

				CRuleCross new_rule;
				new_rule.before = GetSubtreeStr( t, n );
				new_rule.before_cross = GetNumCross( t2, n2,s );
				new_rule.freq = 1;

				int tmp;

				// reorder type1: 0 1 2 -> 0 2 1
				tmp = n2.child[1]; n2.child[1] = n2.child[2]; n2.child[2] = tmp;
				new_rule.after = "A0 A2 A1";
				new_rule.after_cross = GetNumCross( t2, n2, s );
				rule_cross.Add( new_rule );
				RecoverReorder( n2, n );

				// reorder type2: 0 1 2  -> 1 0 2
				tmp = n2.child[0]; n2.child[0] = n2.child[1]; n2.child[1] = tmp;
				new_rule.after = "A1 A0 A2";
				new_rule.after_cross = GetNumCross( t2, n2, s );
				rule_cross.Add( new_rule );
				RecoverReorder( n2, n );

				// reorder type3: 0 1 2  -> 2 0 1
				tmp = n2.child[2]; n2.child[2] = n2.child[1]; n2.child[1] = n2.child[0]; n2.child[0] = tmp;
				new_rule.after = "A2 A0 A1";
				new_rule.after_cross = GetNumCross( t2, n2, s );
				rule_cross.Add( new_rule );
				RecoverReorder( n2, n );

				// reorder type4: 0 1 2  -> 1 2 0
				tmp = n2.child[0]; n2.child[0] = n2.child[1]; n2.child[1] = n2.child[2]; n2.child[2] = tmp;
				new_rule.after = "A1 A2 A0";
				new_rule.after_cross = GetNumCross( t2, n2, s );
				rule_cross.Add( new_rule );
				RecoverReorder( n2, n );

				// reorder type4: 0 1 2  -> 2 1 0
				tmp = n2.child[0]; n2.child[0] = n2.child[2]; n2.child[2] = tmp;
				new_rule.after = "A2 A1 A0";
				new_rule.after_cross = GetNumCross( t2, n2, s );
				rule_cross.Add( new_rule );
				RecoverReorder( n2, n );
			}
		}
	}
	fprintf( stderr, "\r[Extract Rule: CR   ]   %d / %d sentences processed (# Rules: %d)\n", num_sent, num_sent, rule_cross.GetNumRules() );

	rule_cross.SaveToFile( out_files.file_rule_cr.GetFP(), stree );

	return true;
}



//////////////////////////////////////////////////////////////////////////////
bool CCopycatTrain::ExtractRuleRF()
{
#ifdef DEBUG
	fprintf( stdout, "ExtractRuleRF\n" ); fflush( stdout );
#endif

	for( int i=0; i<num_sent; i++ )	{
		fprintf( stderr, "\r[Extract Rule: RF   ]   %d / %d sentences processed (# Rules: %d)", i, num_sent, rule_rf.GetNumRules() );

		CSent &s = sent[i];

		// load tree
		CTree *t = new CTree;
		t->Init();
		if( t->LoadTree( s.parse ) == false )	{
			continue;
			//return false;
		}

		CTree *t2 = new CTree;
		t2->Init();
		if( t2->LoadTree( s.parse ) == false )	{
			continue;
			//return false;
		}


		SetNodeSpan( s, *t );
		SetNodeSpan( s, *t2 );

		// store src statics
		for( int j=0; j<t->num_node; j++ )	{
			CNode &n = t->node[j];
			CNode &n2 = t2->node[j];

			if( MCH == "2" )	{ if( n.num_child != 2 )	continue; }
			else if( MCH == "23" )	 { if( n.num_child == 2  || n.num_child ==3 )	{ } else	{ continue; } }
			else { if( n.num_child <= 1 ) continue;  }

			if( IsAllChildrenAligned(*t, n ) == false ) continue;

			CRuleRF new_rule;
			new_rule.before = GetSubtreeStr( *t, n );
			new_rule.after =  GetSortedSubtreeStr( *t, n );
			new_rule.freq = 1;
			rule_rf.Add( new_rule );
			RecoverReorder( n2, n );
		}

		delete t;
		delete t2;

	}
	fprintf( stderr, "\r[Extract Rule: RF   ]   %d / %d sentences processed (# Rules: %d)\n", num_sent, num_sent, rule_rf.GetNumRules() );

	rule_rf.SaveToFile( out_files.file_rule_rf.GetFP(), stree, false );

#ifdef DEBUG
	fprintf( stdout, "ExtractRuleRF Finish\n" ); fflush( stdout );
#endif
	return true;
}


CUString CCopycatTrain::InsertMissedWords( CTree &t, CNode &n, CSent &s )
{
	CUString ret;

	for( int i=0; i<n.num_child-1; i++ )	{
		int c1id = n.child[i];
		CNode &c1_node = t.node[ c1id ];
		int c2id = n.child[i+1];
		CNode &c2_node = t.node[ c2id ];


		if( c2_node.span_from - c1_node.span_to + 1 >= 4)	{
			return CUString( "ERROR" );
		}

		// print child nodes
		char tmp[128]; sprintf( tmp, " A%d", i ); 
		ret += tmp;

		// print missed words
		for( int j=c1_node.span_to+1; j<c2_node.span_from; j++ )	{
			ret +=  CUString( " I") + s.tgt_tok_list.GetAt(j);
		}
	}

	// print the last child node
	/*
	int i = n.num_child-1;
	int cid = n.child[i];
	CNode &last_node = t.node[ cid ];
	ret += CUString(" [")+last_node.label+CUString("]");
	*/
	char tmp[128]; sprintf( tmp, " A%d", n.num_child-1 ); 
	ret += tmp;

	return ret;
}

//////////////////////////////////////////////////////////////////////////////
bool CCopycatTrain::ExtractInsRule()
{

	for( int i=0; i<num_sent; i++ )	{
		fprintf( stderr, "\r[Extract Ins Rule ]   %d / %d sentences processed (# Rules: %d)", i, num_sent, ins_rule.GetNumRules() );

		CSent &s = sent[i];

		// load tree
		CTree t, t2;
		t.Init(); t2.Init();
		if( t.LoadTree( s.parse ) == false )	{
			continue;
			//return false;
		}
		if( t2.LoadTree( s.parse ) == false )	{
			continue;
			//return false;
		}
		SetNodeSpan( s, t );
		SetNodeSpan( s, t2 );

		// store src statics
		for( int j=0; j<t.num_node; j++ )	{
			CNode &n = t.node[j];
			CNode &n2 = t2.node[j];

			if( MCH == "2" )	{ if( n.num_child != 2 )	continue; }
			else if( MCH == "23" )	 { if( n.num_child == 2  || n.num_child ==3 )	{ } else	{ continue; } }
			else { if( n.num_child <= 1 ) continue;  }

			if( IsAllChildrenAligned( t, n ) == false ) continue;

			CInsertRule new_rule;
			new_rule.before = GetSubtreeStr( t, n );
			new_rule.after = InsertMissedWords( t, n, s );
			if( new_rule.after != "ERROR" )	{
				new_rule.freq = 1;
				ins_rule.Add( new_rule );
				//RecoverReorder( n2, n );
			}
		}

	}
	fprintf( stderr, "\r[Extract Ins Rule ]   %d / %d sentences processed (# Rules: %d)\n", num_sent, num_sent, ins_rule.GetNumRules() );

	ins_rule.MarkBestRule();

	ins_rule.SaveToFile( out_files.file_ins_rule.GetFP(), stree, false );

	return true;
}




